feat: support seeding model from Huggingface repo #81

rickzx · 2024-08-13T09:04:47Z

Example model:

kubectl create secret generic hf-secret --from-literal=HF_TOKEN=YOUR_TOKEN'

  downloaderContainerEnvFrom:
  - secretRef:
      name: hf-secret
  models:
  - downloadUrl: hf://google/gemma-2-2b-it@https://hf-mirror.com
    tag: google--gemma-2-2b-it:e48216d9004e7fd70bc4fdfdc5b7cc3349f8e619

Log output:

❯ kubectl logs yatai-model-seeder-cqthnj0nhgvd725lqp70-mz474
HF_TOKEN is set, using HF_TOKEN to download model...
Downloading model google/gemma-2-2b-it (endpoint=https://hf-mirror.com, revision=e48216d9004e7fd70bc4fdfdc5b7cc3349f8e619) from Huggingface...

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]Downloading 'model-00002-of-00002.safetensors' to '/tmp/model/.cache/huggingface/download/model-00002-of-00002.safetensors.6d6d9ce84db398fb6e0191f91542e5da0a73da2cb695e172a24edc2146dc8d20.incomplete'
Downloading 'model.safetensors.index.json' to '/tmp/model/.cache/huggingface/download/model.safetensors.index.json.022daff51a2616bf88f5f0d0eac5dbd525e7cd24.incomplete'
Downloading 'model-00001-of-00002.safetensors' to '/tmp/model/.cache/huggingface/download/model-00001-of-00002.safetensors.532d792c9178805064170a3ec485b7dedbfccc6fd297b92c31a6091b6c7e41bf.incomplete'
Downloading '.gitattributes' to '/tmp/model/.cache/huggingface/download/.gitattributes.52373fe24473b1aa44333d318f578ae6bf04b49b.incomplete'
Downloading 'special_tokens_map.json' to '/tmp/model/.cache/huggingface/download/special_tokens_map.json.8d6368f7e735fbe4781bf6e956b7c6ad0586df80.incomplete'
Downloading 'generation_config.json' to '/tmp/model/.cache/huggingface/download/generation_config.json.2da7d771ff9ea5f81a30bced0e7390b52061c08e.incomplete'
Downloading 'README.md' to '/tmp/model/.cache/huggingface/download/README.md.f24e534229e1b9e4dd60acc96b2583eb464083f9.incomplete'
Downloading 'config.json' to '/tmp/model/.cache/huggingface/download/config.json.05131f6b339647ddff99327b199ac3d34a50bf2e.incomplete'
Download complete. Moving file to /tmp/model/model.safetensors.index.json
Download complete. Moving file to /tmp/model/special_tokens_map.json
Download complete. Moving file to /tmp/model/.gitattributes
Fetching 11 files:   9%|▉         | 1/11 [00:00<00:09,  1.09it/s]Download complete. Moving file to /tmp/model/generation_config.json
Download complete. Moving file to /tmp/model/README.md
Download complete. Moving file to /tmp/model/config.json
Fetching 11 files:  27%|██▋       | 3/11 [00:01<00:02,  3.56it/s]Downloading 'tokenizer.json' to '/tmp/model/.cache/huggingface/download/tokenizer.json.3f289bc05132635a8bc7aca7aa21255efd5e18f3710f43e3cdb96bcd41be4922.incomplete'
Downloading 'tokenizer.model' to '/tmp/model/.cache/huggingface/download/tokenizer.model.61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2.incomplete'
Downloading 'tokenizer_config.json' to '/tmp/model/.cache/huggingface/download/tokenizer_config.json.36ba12877e7efaba5252f7b17ab337a59d953563.incomplete'
Download complete. Moving file to /tmp/model/tokenizer.model
Download complete. Moving file to /tmp/model/tokenizer_config.json
Download complete. Moving file to /tmp/model/tokenizer.json
Fetching 11 files:  27%|██▋       | 3/11 [00:20<00:02,  3.56it/s]Download complete. Moving file to /tmp/model/model-00001-of-00002.safetensors
Fetching 11 files:  45%|████▌     | 5/11 [00:36<00:56,  9.38s/it]Download complete. Moving file to /tmp/model/model-00002-of-00002.safetensors
Fetching 11 files: 100%|██████████| 11/11 [06:11<00:00, 33.80s/it]
/tmp/model
Moving model to /juicefs-workspace...
Creating /juicefs-workspace/.exists file...
Done

apis/resources/v1alpha1/bentorequest_types.go

controllers/resources/bentorequest_controller.go

yetone · 2024-08-13T09:25:32Z

controllers/resources/bentorequest_controller.go

+	echo "Downloading model {{.ModelRepositoryName}}:{{.ModelVersion}} to /tmp/downloaded.tar..."
+	if [[ ${url} == s3://* ]]; then
+		echo "Downloading from s3..."
+		aws s3 cp ${url} /tmp/downloaded.tar


‌‌‌‌‌‌‌Is this acceptable? This way, there's no need to clean up /tmp/downloaded.tar, it reduces two disk I/O operations, and it decreases the writing to the container's writable layer, thus reducing the usage of ephemeral storage.

aws s3 cp {url} - | tar -xv -C {{.ModelDirPath}}

controllers/resources/bentorequest_controller.go